// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <turbo/log/logging.h>
#include <atomic>
#include <utility>

#include <pollux/common/io/io_statistics.h>

namespace kumo::pollux::io {
    uint64_t IoStatistics::rawBytesRead() const {
        return rawBytesRead_.load(std::memory_order_relaxed);
    }

    uint64_t IoStatistics::rawOverreadBytes() const {
        return rawOverreadBytes_.load(std::memory_order_relaxed);
    }

    uint64_t IoStatistics::rawBytesWritten() const {
        return rawBytesWritten_.load(std::memory_order_relaxed);
    }

    uint64_t IoStatistics::inputBatchSize() const {
        return inputBatchSize_.load(std::memory_order_relaxed);
    }

    uint64_t IoStatistics::outputBatchSize() const {
        return outputBatchSize_.load(std::memory_order_relaxed);
    }

    uint64_t IoStatistics::totalScanTime() const {
        return totalScanTime_.load(std::memory_order_relaxed);
    }

    uint64_t IoStatistics::writeIOTimeUs() const {
        return writeIOTimeUs_.load(std::memory_order_relaxed);
    }

    uint64_t IoStatistics::incRawBytesRead(int64_t v) {
        return rawBytesRead_.fetch_add(v, std::memory_order_relaxed);
    }

    uint64_t IoStatistics::incRawBytesWritten(int64_t v) {
        return rawBytesWritten_.fetch_add(v, std::memory_order_relaxed);
    }

    uint64_t IoStatistics::incInputBatchSize(int64_t v) {
        return inputBatchSize_.fetch_add(v, std::memory_order_relaxed);
    }

    uint64_t IoStatistics::incOutputBatchSize(int64_t v) {
        return outputBatchSize_.fetch_add(v, std::memory_order_relaxed);
    }

    uint64_t IoStatistics::incRawOverreadBytes(int64_t v) {
        return rawOverreadBytes_.fetch_add(v, std::memory_order_relaxed);
    }

    uint64_t IoStatistics::incTotalScanTime(int64_t v) {
        return totalScanTime_.fetch_add(v, std::memory_order_relaxed);
    }

    uint64_t IoStatistics::incWriteIOTimeUs(int64_t v) {
        return writeIOTimeUs_.fetch_add(v, std::memory_order_relaxed);
    }

    void IoStatistics::incOperationCounters(
        const std::string &operation,
        const uint64_t resourceThrottleCount,
        const uint64_t localThrottleCount,
        const uint64_t networkThrottleCount,
        const uint64_t globalThrottleCount,
        const uint64_t retryCount,
        const uint64_t latencyInMs,
        const uint64_t delayInjectedInSecs,
        const uint64_t fullThrottleCount,
        const uint64_t partialThrottleCount) {
        std::lock_guard<std::mutex> lock{operationStatsMutex_};
        operationStats_[operation].localThrottleCount += localThrottleCount;
        operationStats_[operation].resourceThrottleCount += resourceThrottleCount;
        operationStats_[operation].networkThrottleCount += networkThrottleCount;
        operationStats_[operation].globalThrottleCount += globalThrottleCount;
        operationStats_[operation].retryCount += retryCount;
        operationStats_[operation].latencyInMs += latencyInMs;
        operationStats_[operation].requestCount++;
        operationStats_[operation].delayInjectedInSecs += delayInjectedInSecs;
        operationStats_[operation].fullThrottleCount += fullThrottleCount;
        operationStats_[operation].partialThrottleCount += partialThrottleCount;
    }

    melon::F14FastMap<std::string, OperationCounters>
    IoStatistics::operationStats() const {
        std::lock_guard<std::mutex> lock{operationStatsMutex_};
        return operationStats_;
    }

    void IoStatistics::merge(const IoStatistics &other) {
        rawBytesRead_ += other.rawBytesRead_;
        rawBytesWritten_ += other.rawBytesWritten_;
        totalScanTime_ += other.totalScanTime_;

        rawOverreadBytes_ += other.rawOverreadBytes_;
        prefetch_.merge(other.prefetch_);
        read_.merge(other.read_);
        ramHit_.merge(other.ramHit_);
        ssdRead_.merge(other.ssdRead_);
        queryThreadIoLatency_.merge(other.queryThreadIoLatency_); {
            const auto &otherOperationStats = other.operationStats();
            std::lock_guard<std::mutex> l(operationStatsMutex_);
            for (auto &item: otherOperationStats) {
                operationStats_[item.first].merge(item.second);
            }
        }
    }

    void OperationCounters::merge(const OperationCounters &other) {
        resourceThrottleCount += other.resourceThrottleCount;
        localThrottleCount += other.localThrottleCount;
        networkThrottleCount += other.networkThrottleCount;
        globalThrottleCount += other.globalThrottleCount;
        retryCount += other.retryCount;
        latencyInMs += other.latencyInMs;
        requestCount += other.requestCount;
        delayInjectedInSecs += other.delayInjectedInSecs;
        fullThrottleCount += other.fullThrottleCount;
        partialThrottleCount += other.partialThrottleCount;
    }

    melon::Dynamic serialize(const OperationCounters &counters) {
        melon::Dynamic json = melon::Dynamic::object;
        json["latencyInMs"] = counters.latencyInMs;
        json["localThrottleCount"] = counters.localThrottleCount;
        json["networkThrottleCount"] = counters.networkThrottleCount;
        json["resourceThrottleCount"] = counters.resourceThrottleCount;
        json["globalThrottleCount"] = counters.globalThrottleCount;
        json["retryCount"] = counters.retryCount;
        json["requestCount"] = counters.requestCount;
        json["delayInjectedInSecs"] = counters.delayInjectedInSecs;
        json["fullThrottleCount"] = counters.fullThrottleCount;
        json["partialThrottleCount"] = counters.partialThrottleCount;
        return json;
    }

    melon::Dynamic IoStatistics::getOperationStatsSnapshot() const {
        auto snapshot = operationStats();
        melon::Dynamic json = melon::Dynamic::object;
        for (auto stat: snapshot) {
            json[stat.first] = serialize(stat.second);
        }
        return json;
    }
} // namespace kumo::pollux::io
